import math
import torch


def qmul(q, r):
    """
    Multiply quaternion(s) q with quaternion(s) r.
    Expects two equally-sized tensors of shape (*, 4), where * denotes any number of dimensions.
    Returns q*r as a tensor of shape (*, 4).
    """
    assert q.shape[-1] == 4
    assert r.shape[-1] == 4

    # Compute outer product
    terms = r.unsqueeze(-1) @ q.unsqueeze(-2)

    w = terms[..., 0, 0] - terms[..., 1, 1] - terms[..., 2, 2] - terms[..., 3, 3]
    x = terms[..., 0, 1] + terms[..., 1, 0] - terms[..., 2, 3] + terms[..., 3, 2]
    y = terms[..., 0, 2] + terms[..., 1, 3] + terms[..., 2, 0] - terms[..., 3, 1]
    z = terms[..., 0, 3] - terms[..., 1, 2] + terms[..., 2, 1] + terms[..., 3, 0]
    return torch.stack((w, x, y, z), dim=-1)


def qrot(q, v):
    """
    Rotate vector(s) v about the rotation described by quaternion(s) q.
    Expects a tensor of shape (*, 4) for q and a tensor of shape (*, 3) for v,
    where * denotes any number of dimensions.
    Returns a tensor of shape (*, 3).
    """
    assert q.shape[-1] == 4
    assert v.shape[-1] == 3
    assert q.shape[:-1] == v.shape[:-1]

    qvec = q[..., 1:]
    # uv = torch.cross(qvec, v, dim=-1)
    # uuv = torch.cross(qvec, uv, dim=-1)
    uv = torch.cross(qvec.double(), v.double(), dim=-1)
    uuv = torch.cross(qvec.float(), uv.float(), dim=-1)
    return v + 2 * (q[..., :1] * uv + uuv)


def qinv(q):
    assert q.shape[-1] == 4
    return torch.tensor([1, -1, -1, -1], dtype=torch.float32, device=q.device) * q
